package tool

import (
	"bytes"
	"crypto/cipher"
	"errors"
	"fmt"
	"github.com/tjfoc/gmsm/sm4"
)

// SM4CBCEncrypt SM4 CBC模式加密
func SM4CBCEncrypt(key, iv, plaintext []byte) ([]byte, error) {
	// 参数检查
	if len(key) != sm4.BlockSize {
		return nil, fmt.Errorf("SM4密钥长度必须为%d字节", sm4.BlockSize)
	}
	if len(iv) != sm4.BlockSize {
		return nil, fmt.Errorf("IV长度必须为%d字节", sm4.BlockSize)
	}

	// 创建加密块
	block, err := sm4.NewCipher(key)
	if err != nil {
		return nil, fmt.Errorf("创建SM4加密块失败: %v", err)
	}

	// 填充数据
	paddedText := PKCS7Padding(plaintext, sm4.BlockSize)

	// CBC模式加密
	mode := cipher.NewCBCEncrypter(block, iv)
	ciphertext := make([]byte, len(paddedText))
	mode.CryptBlocks(ciphertext, paddedText)

	return ciphertext, nil
}

// SM4CBCDecrypt SM4 CBC模式解密
// 如果原字符串是 base64 或者 hex 需要先解码
//
//	 hexString, err := hex.DecodeString("a37c8266d25b9e761ca5f7c5d22ab75ab5705836ab0c9ab4c93b20f15f9364d442fcf673e8aab245d2eddab8429e07b2f3ca12f45deccdea82db490aa0c727914fd1f349d648d203cdc51a953643c1cbe5f2e7984aa0ec425e111249364d4f0a9fde30bfea093268f8895d5dbdb48c35683eb016529fdeb48d627ec67b686b3e4a421e78573ec6b869d0b5a1d030f2b025e45603a532cc5e8aed04e3afe02c3df19b0a7575584a7cce43240628566dc0e600df8911baefbaaff9f99cc500b935dd03e00fad059afc69d45c86c44862029ddeb2928f81335f7569b253abd8981e1a6d63eb044ba21bb8eeb5467ec452cf2e57c0b0ff35c96757c36753fe4aec2639a612eacce37e251e022346c44441e2")
//		decrypted, err := tool.SM4CBCDecrypt(key, iv, hexString)
func SM4CBCDecrypt(key, iv, ciphertext []byte) ([]byte, error) {
	// 参数检查
	if len(key) != sm4.BlockSize {
		return nil, fmt.Errorf("SM4密钥长度必须为%d字节", sm4.BlockSize)
	}
	if len(iv) != sm4.BlockSize {
		return nil, fmt.Errorf("IV长度必须为%d字节", sm4.BlockSize)
	}
	if len(ciphertext)%sm4.BlockSize != 0 {
		return nil, fmt.Errorf("密文长度必须是%d的倍数", sm4.BlockSize)
	}

	// 创建解密块
	block, err := sm4.NewCipher(key)
	if err != nil {
		return nil, fmt.Errorf("创建SM4解密块失败: %v", err)
	}

	// CBC模式解密
	mode := cipher.NewCBCDecrypter(block, iv)
	plaintext := make([]byte, len(ciphertext))
	mode.CryptBlocks(plaintext, ciphertext)

	// 去除填充
	plaintext, err = PKCS7UnPadding(plaintext)
	if err != nil {
		return nil, fmt.Errorf("去除填充失败: %v", err)
	}

	return plaintext, nil
}

// PKCS7Padding 数据填充
func PKCS7Padding(data []byte, blockSize int) []byte {
	padding := blockSize - len(data)%blockSize
	padText := bytes.Repeat([]byte{byte(padding)}, padding)
	return append(data, padText...)
}

// PKCS7UnPadding 去除填充
func PKCS7UnPadding(data []byte) ([]byte, error) {
	if len(data) == 0 {
		return nil, errors.New("输入数据为空")
	}
	padding := int(data[len(data)-1])
	if padding < 1 || padding > sm4.BlockSize {
		return nil, errors.New("填充大小无效")
	}
	if len(data) < padding {
		return nil, errors.New("数据长度小于填充大小")
	}
	return data[:len(data)-padding], nil
}
